from AStock.ASEvent import Task, Engine
from AStock.ASCrawler import (
    ASFinancialNoticeDateSaver,
    ASFinancialNoticeDateCrawler,
    ASFinancialSaver,
    ASFinancialCrawler,
    ASFinancialPreviewCrawler,
    ASFinancialPreviewSaver
)
from AStock.ASSetting import settings  # , cache_path
# from AStock.ASUtil import LRUCache
# from datetime import datetime
import random
import time
import argparse
# import os


class ASSaveFinancial(object):

    def __init__(self, report_date, stocks, engine=None):
        self._report_date = report_date
        self._stocks = stocks
        if engine is None:
            self._engine = Engine()
            self._engine.start()
            self._internal_engine = True
        else:
            self._engine = engine
            self._internal_engine = False

    def run(
            self,
            financial=True,
            financial_preview=False
    ):
        if financial_preview:
            self._run_financial_preview()

        if financial:
            # 处理获取报表，包括报告日期、资产负债表、现金流量表、利润表、财务比率
            self._run_financial()

        if self._internal_engine:
            # 等待任务完成
            self._engine.safe_quit()
            print('done')

    def _crawl_one_stock_financial(self, code):
        print('>>>>>>>>>>>> {} save financial >>>>>>>>>>>>'.format(code))
        event_types = ['zcfzb', 'xjllb', 'lrb', 'cwbl']
        for event_type in event_types:
            event_data = {
                'event_type': event_type,
                'report_date': self._report_date,
                'code': code
            }

            task = Task(
                event_data,
                ASFinancialCrawler(),
                bind_event='financial_' + event_type,
                callback=self._financial_result_handler
            )

            self._engine.run_task(task)
            time.sleep(0.3)

    def _financial_notice_date_result_handler(self, result, **kwargs):
        """
        在这里保存抓取的财报发布日期，并抓取下一页
        :param result: 抓取的财报发布日期结果数据，
                        {'error': 'error string'}
                        {'page_no': 1, 'pages': 105, 'data': [(code, notice_date), (code, notice_date), ...]}
        :param kwargs:
        :return:
        """
        if 'error' in result:
            print(result)
        else:
            try:
                rand = kwargs['rand']
                page_no = kwargs['page_no']
            except KeyError as e:
                print('error process financial notice date: {}  key error'.format(e))
                return
            print({'succ': 'crawled   {}    financial notice date'.format(self._report_date)})

            # 生成保存数据到mongodb的任务
            data = [r for r in result['data'] if r[0] in self._stocks] if self._stocks else result['data']
            event_type = 'financial_notice_date'
            event_data = {
                'event_type': event_type,
                'report_date': self._report_date,
                'data': data
            }
            if event_data['data']:
                task = Task(
                    event_data,
                    ASFinancialNoticeDateSaver(),
                    bind_event='financial_notice_date_saver',
                    callback=print
                )

                self._engine.run_task(task)

            # 启动抓取财务报表任务
            for d in data:
                self._crawl_one_stock_financial(d[0])

            # 继续抓取下一页
            if page_no < result['pages']:
                event_data = {
                    'event_type': event_type,
                    'report_date': self._report_date,
                    'page_no': page_no+1,
                    'page_size': 50,
                    'rand': rand+1
                }

                task = Task(
                    event_data,
                    ASFinancialNoticeDateCrawler(),
                    bind_event='financial_notice_date_crawler',
                    callback=self._financial_notice_date_result_handler,
                    page_no=event_data['page_no'],
                    rand=event_data['rand']
                )

                self._engine.run_task(task)

    def _run_financial_notice_date(self):
        """
        触发抓取财报公告日期（第一页）
        :param self._engine:
        :return:
        """
        # 处理发布日期
        self._engine.create_event('financial_notice_date_crawler')
        self._engine.start_event('financial_notice_date_crawler')
        self._engine.create_event('financial_notice_date_saver')
        self._engine.start_event('financial_notice_date_saver')

        event_type = 'financial_notice_date'
        event_data = {
            'event_type': event_type,
            'report_date': self._report_date,
            'page_no': 1,
            'page_size': 50,
            'rand': random.randint(50000000, 57000000)
        }

        task = Task(
            event_data,
            ASFinancialNoticeDateCrawler(),
            bind_event='financial_notice_date_crawler',
            callback=self._financial_notice_date_result_handler,
            page_no=event_data['page_no'],
            rand=event_data['rand']
        )

        self._engine.run_task(task)

    def _filter_stocks_for_financial(self, report_date, force_re_save=False):
        """
        根据当前日期及financial表中“580财报公告日期”字段的时间
        获取已发布财报的股票代码（废弃）
        根据@force_re_save参数需要获取财报的股票代码列表进行过滤
        :param report_date: str, report date
        :param force_re_save: boolean, 库中已经有数据时，是否强制重新获取并保存数据
        :return: list, stock code list to save financial report
        """
        if force_re_save:
            return self._stocks

        # date = datetime.now().strftime('%Y%m%d')
        # date = int(date)

        # filter = {'report_date': int(report_date), '580': {'$lt': date}}

        # 非空则表示已经获取，不需要再保存：

        # 073负债和所有者（或股东权益）合计

        # 096归属于母公司所有者的净利润

        # 101经营活动现金流入小计

        # 197净资产收益率
        # 002扣除非经常性损益每股收益

        filters = {
            'report_date': int(report_date),
            '073': {'$gte': 0.0},
            '096': {'$gte': 0.0},
            '101': {'$gte': 0.0},
            '197': {'$gte': 0.0},
            '002': {'$gte': 0.0}
        }

        """
        filter['$or'] = [
            {'073': 0.0}, {'073': None},
            {'096': 0.0}, {'096': None},
            {'101': 0.0}, {'101': None},
            {'281': 0.0}, {'281': None},
            {'002': 0.0}, {'002': None}
        ]
        """

        coll = settings.database.financial
        res = []
        try:
            excludes = coll.find(filters, {'code': 1})
            exclude_set = set()
            for item in excludes:
                exclude_set.add(item['code'])
            for stock in self._stocks:
                if stock not in exclude_set:
                    res.append(stock)
                else:
                    print('== {} financial already saved, skip it =='.format(stock))
        except Exception as e:
            print('failed to find stocks with financial already saved, {}'.format(e))
            res = self._stocks
        return res

    def _financial_result_handler(self, result, **kwargs):
        """
        在这里保存抓取的财报
        :param result:
        :param kwargs:
        :return:
        """
        if 'error' in result:
            print(result)
            return

        print({'succ': 'crawled   {code}    {report_date}    '
                       'financial {event_type}'.format(**result)})

        task = Task(
            result,
            ASFinancialSaver(),
            bind_event='financial_' + result.get('event_type'),
            callback=self._financial_save_result_handler
        )
        self._engine.run_task(task)

    @staticmethod
    def _financial_save_result_handler(result, **kwargs):
        """
        将已抓取的code更新到cache
        :param kwargs:
        :return:
        """
        print(result)

    def _run_financial(self):
        """
        获取并保存财务报表
        :param self._engine: 事件引擎
        :param force_re_save: 库中已经有数据时，是否强制重新获取并保存数据
        :return:
        """
        # 按财务报表类别创建事件线程，爬取数据和保存数据共用一个事件线程
        self._engine.create_event('financial_zcfzb')
        self._engine.create_event('financial_xjllb')
        self._engine.create_event('financial_cwbl')
        self._engine.create_event('financial_lrb')

        self._engine.start_event('financial_xjllb')
        self._engine.start_event('financial_zcfzb')
        self._engine.start_event('financial_lrb')
        self._engine.start_event('financial_cwbl')

        # 启动报表日期抓取，并在回调中抓取财务报表
        self._run_financial_notice_date()

    @classmethod
    def _financial_preview_result_handler(cls, result, **kwargs):
        """
        在这里保存抓取的业绩预告/快报，并抓取下一页
        :param result: 抓取的业绩预告/快报结果数据，
                        {'error': 'error string'}
                        {'page_no': 1, 'pages': 105, 'data': [..., ...]}
                其中data元素为业绩预告/快报数据：
                {
                    'code': '300001',
                    'notice_date': '20200821',
                    'forcast_content': '预计2020年1-3月归属于上市公司股东的净利润盈利:约6,680万元-8,164万元,同比增长约22.45%-49.66%。',
                    'change_reason': '2020年1-3月经营业绩有所下降,主要系受疫情影响',
                    'forcast_type': '略减',
                    'forcast_net_profit': 11545000,  # 净利（扣非）
                    'forcast_income': 741553800  # 营收
                }
        :param kwargs:
        :return:
        """
        if 'error' in result:
            print(result)
            return

        try:
            engine = kwargs['engine']
            rand = kwargs['rand']
            page_no = kwargs['page_no']
            report_date = kwargs['report_date']
            stocks = kwargs.get('stocks', None)  # 股票code列表，只保存指定的股票
            event_type = kwargs['event_type']
        except KeyError as e:
            print('error process financial preview: {} key error'.format(e))
            return
        print({'succ': 'crawled   {}    financial preview {}'.format(report_date, event_type)})

        # for r in result['data']:
        #     if r['code'] not in stocks:
        #         print("======== {} ========".format(r['code']))

        # 生成保存数据到mongodb的任务
        event_data = {
            'event_type': event_type,
            'report_date': report_date,
            'data': [r for r in result['data'] if r['code'] in stocks] if stocks else result['data']
        }
        if event_data['data']:
            task = Task(
                event_data,
                ASFinancialPreviewSaver(),
                bind_event='financial_preview_saver',
                callback=print
            )

            engine.run_task(task)

        # 继续抓取下一页
        if page_no < result['pages']:
            event_data = {
                'event_type': event_type,
                'report_date': report_date,
                'page_no': page_no+1,
                'page_size': 100,
                'rand': rand+1
            }

            task = Task(
                event_data,
                ASFinancialPreviewCrawler(),
                bind_event='financial_preview_crawler',
                callback=cls._financial_preview_result_handler,
                page_no=event_data['page_no'],
                rand=event_data['rand'],
                report_date=event_data['report_date'],
                engine=engine,
                stocks=stocks,
                event_type=event_type
            )

            engine.run_task(task)

    def _run_financial_preview(self):
        """
        触发抓取业绩预告/业绩快报（第一页）
        :param self._engine:
        :return:
        """
        # 处理发布日期
        self._engine.create_event('financial_preview_crawler')
        self._engine.start_event('financial_preview_crawler')
        self._engine.create_event('financial_preview_saver')
        self._engine.start_event('financial_preview_saver')

        event_type = ['yjyg', 'yjkb']
        for et in event_type:
            event_data = {
                'event_type': et,
                'report_date': self._report_date,
                'page_no': 1,
                'page_size': 100,
                'rand': random.randint(50000000, 57000000)
            }

            task = Task(
                event_data,
                ASFinancialPreviewCrawler(),
                bind_event='financial_preview_crawler',
                callback=self._financial_preview_result_handler,
                page_no=event_data['page_no'],
                rand=event_data['rand'],
                report_date=event_data['report_date'],
                engine=self._engine,
                stocks=self._stocks,
                event_type=et
            )

            self._engine.run_task(task)


def remove_exists(report_date, stocks, financial, financial_preview):
    """
    移除数据库里已经存在财务数据的股票
    数据库里是否存在财务数据的判断依据：
    正式财报：‘580财报公告日期’ ‘206扣除非经常性损益后的净利润’ ‘074一、营业总收入’ 数据均存在
    业绩预告：‘578预计净利润’ 数据存在 且 ‘581业绩快报公告日期’ 或 ‘582业绩预告公告日期’ 数据存在
    :param report_date: str yyyymmdd 报告期
    :param stocks: list of str 待过滤的股票列表
    :param financial: bool 是否正式财报
    :param financial_preview: bool 是否业绩预告
    :return: list of str 过滤后剩余的股票列表
    """
    filter_ = {'report_date': int(report_date)}
    if financial:
        filter_.__setitem__('580', {'$ne': None})
        filter_.__setitem__('206', {'$ne': None})
        filter_.__setitem__('074', {'$ne': None})
    if financial_preview:
        filter_.__setitem__('578', {'$ne': None})
        filter_.__setitem__('$or', {'581': {'$ne': None}, '582': {'$ne': None}})
    coll = settings.database.financial
    # print(filter_)
    try:
        cursor = coll.find(filter_, {'code': 1})
    except Exception as e:
        print(e)
        return stocks
    to_be_removed = set([row['code'] for row in cursor])
    print('{} already in database'.format(len(to_be_removed)))
    res = [s for s in stocks if s not in to_be_removed]
    return res


def _main_(args):
    report_date = args.report_date
    financial = args.financial
    stocks = args.stock_codes
    financial_preview = args.financial_preview
    skip_exists = args.skip_exists

    print('report_date: {}, financial: {}, financial_preview: {}, skip_exists: {}'.format(
        report_date, financial, financial_preview, skip_exists))

    # print(settings)

    if not stocks:
        # 保存stock_info_tushare集合中存在的stock
        try:
            coll = settings.database.stock_info_tushare
            stocks = coll.find(
                {},
                {'code': 1, '_id': 0}
            )

            stocks = [stock['code'] for stock in stocks]
        except Exception as e:
            print(e)
            exit(1)

    if skip_exists:
        stocks = remove_exists(report_date, stocks, financial, financial_preview)

    sf = ASSaveFinancial(report_date, stocks)
    sf.run(
        financial=financial,
        financial_preview=financial_preview
    )


if __name__ == '__main__':
    argparser = argparse.ArgumentParser(description='save financial report')
    argparser.add_argument('-d', '--report-date', required=True,
                           help='report date of financial report to save, yyyymmdd format')

    flag_parser = argparser.add_mutually_exclusive_group(required=False)
    flag_parser.add_argument('--financial', dest='financial', action='store_true',
                             help='save financial report')
    flag_parser.add_argument('--no-financial', dest='financial', action='store_false',
                             help="don't save financial report")
    argparser.set_defaults(financial=True)
    """
    获取财报前先获取财报发布日期，
    通过财报发布日期获得已经发布财报的股票，据此获取这些股票的财报，避免请求所有股票财报URL
    flag_parser = argparser.add_mutually_exclusive_group(required=False)
    flag_parser.add_argument('--financial-notice-date', dest='financial_notice_date', action='store_true',
                             help='save financial notice date')
    flag_parser.add_argument('--no-financial-notice-date', dest='financial_notice_date', action='store_false',
                             help="don't save financial notice date")
    argparser.set_defaults(financial_notice_date=True)
    """
    flag_parser = argparser.add_mutually_exclusive_group(required=False)
    flag_parser.add_argument('--financial-preview', dest='financial_preview', action='store_true',
                             help='save financial preview')
    flag_parser.add_argument('--no-financial-preview', dest='financial_preview', action='store_false',
                             help="don't save financial preview")
    argparser.set_defaults(financial_preview=False)

    argparser.add_argument('--stock-codes', nargs='+', help='stock codes to save financial')

    argparser.add_argument('--skip-exists', required=False, action='store_true',
                           help='skip exists when saving report')

    args = argparser.parse_args()
    _main_(args)


